#include "ImageHost.h"
#include "Utils.h"
#include "BlockDivider.h"
#include "CvUtil.h"
using namespace cv;

struct opRange
{
    int x1;
    int x2;
};
struct opRange2D
{
    int x1;
    int x2;
    int y1;
    int y2;
    int stepX;
    int stepY;
};
static void gen_rangeyx(int dir, const opRange2D &range, opRange2D &out)
{
    if (dir == 0)
    {
        out.x1 = range.x1;
        out.x2 = range.x2;
        out.y1 = range.y1;
        out.y2 = range.y2;
        out.stepX = 1;
        out.stepY = 1;
    }
    else if (dir == 1)
    {
        out.x1 = range.x2 - 1;
        out.x2 = range.x1 - 1;
        out.y1 = range.y1;
        out.y2 = range.y2;
        out.stepX = -1;
        out.stepY = 1;
    }
    else if (dir == 2)
    {
        out.x1 = range.x1;
        out.x2 = range.x2;
        out.y1 = range.y2 - 1;
        out.y2 = range.y1 - 1;
        out.stepX = 1;
        out.stepY = -1;
    }
    else
    {
        out.x1 = range.x2 - 1;
        out.x2 = range.x1 - 1;
        out.y1 = range.y2 - 1;
        out.y2 = range.y1 - 1;
        out.stepX = -1;
        out.stepY = -1;
    }
}

ImageHost::ImageHost()
    : m_threadPool(std::thread::hardware_concurrency())
{
}

ImageHost::~ImageHost()
{
}

long ImageHost::Capture(const std::wstring &file)
{
    if (_src.empty())
    {
        setlog(L"_src is empty");
        return 0;
    }
    if (!cv::imwrite(_ws2string(file), _src))
    {
        setlog(L"imwrite failed");
        return 0;
    }
    return 1;
}

void ImageHost::set_offset(int x1, int y1)
{
    _x1 = x1;
    _y1 = y1;
}

long ImageHost::CmpColor(long x, long y, const std::wstring &scolor, double sim)
{
    std::vector<color_df_t> vcolor;
    str2colordfs(scolor, vcolor);

    cv::Vec4b c = _src.at<cv::Vec4b>(0, 0);
    color_t color(c[0], c[1], c[2]);

    for (auto &it : vcolor)
    {
        if (IN_RANGE(color, it.color, it.df))
            return 1;
    }

    return 0;
}

long ImageHost::FindColor(const wstring &color, double sim, long dir, long &x, long &y)
{
    std::vector<color_df_t> colors;
    str2colordfs(color, colors);

    opRange2D rng = {0, _src.cols, 0, _src.rows, 0, 0}, range;
    gen_rangeyx(dir, rng, range);

    for (auto &it : colors)
    {
        for (int i = range.y1; i != range.y2; i += range.stepY)
        {
            for (int j = range.x1; j != range.x2; j += range.stepX)
            {
                cv::Vec4b c = _src.at<cv::Vec4b>(i, j);
                color_t color(c[0], c[1], c[2]);
                if (IN_RANGE(color, it.color, it.df))
                {
                    x = j + _x1 + _dx;
                    y = i + _y1 + _dy;
                    return 1;
                }
            }
        }
    }
    x = y = -1;
    return 0;
}

long ImageHost::FindPic(const std::wstring &files, const wstring &delta_colors, double sim, long dir, long &x, long &y)
{
    // TODO
    vector<cv::Mat> vpic;
    // vector<color_t> vcolor;
    color_t dfcolor;
    vector<std::wstring> vpic_name;
    files2mats(files, vpic, vpic_name);
    dfcolor.str2color(delta_colors);
    // str2colors(delta_colors, vcolor);
    sim = 0.5 + sim / 2;

    return 0;
}

long ImageHost::FindPic(const std::wstring &files, double threshold, long &x1, long &y1, long &x2, long &y2)
{
    vector<cv::Mat> pics;
    vector<std::wstring> vpic_name;
    files2mats(files, pics, vpic_name);

    for (int pic_id = 0; pic_id < pics.size(); ++pic_id)
    {
        auto pic = pics[pic_id];

        // cv::Rect rect = findTemplateWithMask(_src, pic, threshold);
        cv::Rect rect = findTemplate(_src, pic, threshold);
        if (rect.x != -1 && rect.y != -1)
        {
            x1 = rect.x + _x1 + _dx;
            y1 = rect.y + _y1 + _dy;
            x2 = x1 + rect.width;
            y2 = y1 + rect.height;
            return 1;
        }
    }
    return 0;
}

long ImageHost::LoadPic(const wstring &files)
{
    std::vector<wstring> vstr;
    int loaded = 0;
    split(files, vstr, L"|");
    wstring tp;
    for (auto &it : vstr)
    {
        // 路径转化
        if (!Path2GlobalPath(it, currPath_, tp))
            continue;
        // 先在缓存中查找
        if (!_pic_cache.count(tp))
        {
            cv::Mat mat = cv::imread(_ws2string(tp));
            if (mat.empty())
            {
                setlog(L"imread failed: %ls", _ws2string(tp).c_str());
                continue;
            }
            _pic_cache[tp] = mat;
        }
        // 已存在于缓存中的文件也算加载成功
        loaded++;
    }
    return loaded;
}

long ImageHost::FreePic(const wstring &files)
{
    std::vector<wstring> vstr;
    int loaded = 0;
    split(files, vstr, L"|");
    wstring tp;
    for (auto &it : vstr)
    {
        // 看当前目录
        auto cache_it = _pic_cache.find(it);
        // 没查到再看一下资源目录
        if (cache_it == _pic_cache.end())
        {
            cache_it = _pic_cache.find(currPath_ + L"\\" + it);
        }
        // 查到了就释放
        if (cache_it != _pic_cache.end())
        {
            _pic_cache.erase(cache_it);
            loaded++;
        }
    }
    return loaded;
}

long ImageHost::LoadMemPic(const wstring &file_name, void *data, long size)
{
    try
    {
        if (!_pic_cache.count(file_name))
        {
            cv::Mat mat = cv::imdecode(cv::Mat(1, size, CV_8UC1, data), cv::IMREAD_UNCHANGED);
            if (mat.empty())
            {
                setlog(L"imdecode failed");
                return 0;
            }
            _pic_cache[file_name] = mat;
        }
    }
    catch (...)
    {
        return 0;
    }
    return 1;
}
long ImageHost::UseDict(const wchar_t *dict_name)
{
    if (!dicts_.count(dict_name))
    {
        Dict dict;
        std::wstring path = currPath_ + L"\\" + dict_name + L".json";
        if (!dict.loadFromFile(path))
        {
            setlog(L"load dict failed: %ls", _ws2string(path).c_str());
            return 0;
        }
        dicts_[dict_name] = dict;
        currDict_ = &dicts_[dict_name];
    }
    else
    {
        currDict_ = &dicts_[dict_name];
    }
    return 1;
}
long ImageHost::ClearDict()
{
    dicts_.clear();
    currDict_ = nullptr;
    return 0;
}
long ImageHost::FindStr(const wstring &strs, double threshold, long &retx1, long &rety1, long &retx2, long &rety2)
{
    vector<cv::Mat> pics;
    std::vector<DictContent> vdict;
    if (!currDict_)
        return 0;
    if (currDict_->dict_info_.mode_ == DictMode::DictMode_SingleWord)
    {
        if (!loadDictsBySingleWord(strs, vdict))
            return 0;
    }
    else if (currDict_->dict_info_.mode_ == DictMode::DictMode_MultiWord)
    {
        // 多字模式
        if (!loadDictsByMultiWord(strs, vdict))
            return 0;
    }
    cv::Mat binSource = getBinaryByCurrDict();
    if (binSource.empty())
        return 0;

    std::vector<std::future<std::vector<cv::Rect>>> futures;

    for (auto &it : vdict)
    {
        auto future = m_threadPool.enqueue([this, threshold](cv::Mat &binSource, cv::Mat &mat)
                                           {
                                 std::vector<cv::Rect> rects;
                                 if (!findDictTemplate(binSource, mat, threshold, rects))
                                 {
                                    //  return 0;
                                 }
                                return rects; }, binSource, it.mat_);
        futures.push_back(std::move(future));
    }

    std::vector<std::vector<cv::Rect>> results;
    for (auto &future : futures)
    {
        results.push_back(future.get());
    }
    cv::Rect rect = CvUtil::findConnectedRectangles(results);
    // cv::Rect rect = findConnectedRectanglesOptimized(results);
    if (rect.x != -1 && rect.y != -1)
    {
        retx1 = rect.x + _x1 + _dx;
        rety1 = rect.y + _y1 + _dy;
        retx2 = rect.x + rect.width + _x1 + _dx;
        rety2 = rect.y + rect.height + _y1 + _dy;
        return 1;
    }
    return 0;
}
std::vector<std::wstring> ImageHost::Ocr(double threshold, std::vector<cv::Rect> &rects)
{

    cv::Mat binSource = getBinaryByCurrDict();
    if (binSource.empty())
        return std::vector<std::wstring>();

    std::vector<std::future<std::vector<OcrResult>>> futures;

    for (auto &it : currDict_->dicts_)
    {
        if (it.second.mat_.empty())
            if (!loadCurrDictMat(it.second.name_))
                continue;
        auto future = m_threadPool.enqueue([this, threshold](cv::Mat &binSource, cv::Mat &mat, const wstring &name)
                                           {
                                 std::vector<cv::Rect> rects;
                                 if (!findDictTemplate(binSource, mat, threshold, rects))
                                 {
                                    //  return 0;
                                 }
                                 std::vector<OcrResult> ors;
                                 for (auto &rect : rects)
                                 {
                                    OcrResult or;
                                    or.name_ = name;
                                    or.rect_ = rect;
                                    ors.push_back(or);
                                 }

                                return ors; }, binSource, it.second.mat_, it.second.name_);
        futures.push_back(std::move(future));
    }

    std::vector<OcrResult> results;
    for (auto &future : futures)
    {
        auto ors = future.get();
        if (ors.empty())
            continue;
        results.insert(results.end(), ors.begin(), ors.end());
    }
    return CvUtil::sortOcrResults(results, rects);
}
bool ImageHost::loadCurrDictMat(const wstring &word_str)
{
    std::wstring path = currPath_ + L"\\" + currDict_->dict_info_.name_ + L"\\";

    if (!currDict_->dicts_.count(word_str))
        return false;
    if (currDict_->dicts_[word_str].mat_.empty())
    {
        path += currDict_->dicts_[word_str].res_name_;
        if (!CvUtil::imread_unicode(path, currDict_->dicts_[word_str].mat_, cv::IMREAD_GRAYSCALE))
            return false;
    }
    return true;
}
int ImageHost::str2colordfs(const std::wstring &color_str, std::vector<color_df_t> &colors)
{
    std::vector<wstring> vstr, vstr2;
    color_df_t cr;
    colors.clear();
    int ret = 0;
    if (color_str.empty())
    { // default
        return 1;
    }
    if (color_str[0] == L'#')
    { // bk color info
        ret = 1;
    }
    split(ret ? color_str.substr(1) : color_str, vstr, L"|");
    for (auto &it : vstr)
    {
        split(it, vstr2, L"-");
        cr.color.str2color(vstr2[0]);
        cr.df.str2color(vstr2.size() == 2 ? vstr2[1] : L"000000");
        colors.push_back(cr);
    }
    return ret;
}

void ImageHost::str2colors(const std::wstring &color, std::vector<color_t> &vcolor)
{
    std::vector<wstring> vstr, vstr2;
    color_t cr;
    vcolor.clear();
    split(color, vstr, L"|");
    for (auto &it : vstr)
    {
        cr.str2color(it);
        vcolor.push_back(cr);
    }
}

void ImageHost::files2mats(const wstring &files, std::vector<cv::Mat> &vpic, std::vector<wstring> &vstr)
{
    cv::Mat pm;
    vpic.clear();
    split(files, vstr, L"|");
    wstring tp;
    for (auto &it : vstr)
    {
        // 先在缓存中查找是否已加载，包括从内存中加载的文件
        if (_pic_cache.count(it))
        {
            pm = _pic_cache[it];
        }
        else
        {
            std::wstring path = currPath_ + L"\\pic";
            // 路径转化
            if (!Path2GlobalPath(it, path, tp))
                continue;
            // 再检测一次，包括绝对路径的文件
            if (_pic_cache.count(tp))
            {
                pm = _pic_cache[tp];
            }
            else
            {
                pm = cv::imread(_ws2string(tp));
                if (pm.empty())
                {
                    setlog(L"imread failed: %ls", _ws2string(tp).c_str());
                    continue;
                }
                _pic_cache[tp] = pm;
            }
        }
        vpic.push_back(pm);
    }
}
cv::Mat ImageHost::getBinaryByCurrDict()
{
    color_t dfcolor;
    if (currDict_ == nullptr)
        return cv::Mat();
    // 去除字符串开头的#字符
    std::wstring colorStr = currDict_->dict_info_.color_;
    if (!colorStr.empty() && colorStr[0] == L'#')
    {
        colorStr = colorStr.substr(1);
    }
    dfcolor.str2color(colorStr);
    cv::Vec3b targetColor(dfcolor.b, dfcolor.g, dfcolor.r);
    int tolerance = _wtoi(currDict_->dict_info_.tolerance_.c_str());

    cv::Mat binSource = CvUtil::colorBinary(_src, targetColor, tolerance);

    return binSource;
}
bool ImageHost::loadDictsBySingleWord(const wstring &strs, std::vector<DictContent> &vdict)
{
    for (const wchar_t &ch : strs)
    {
        wcout << ch << endl;
        wstring temp;
        temp.append(1, ch);
        if (!loadCurrDictMat(temp))
            return false;
        vdict.push_back(currDict_->dicts_[temp]);
    }
    return true;
}
bool ImageHost::loadDictsByMultiWord(const wstring &strs, std::vector<DictContent> &vdict)
{
    std::wstring path = currPath_ + L"\\" + currDict_->dict_info_.name_ + L"\\";
    size_t i = 0;
    size_t len = strs.length();

    // 预先获取所有键并按长度降序排序，这样我们可以优先匹配最长的键
    std::vector<std::wstring> keys;
    for (const auto &pair : currDict_->dicts_)
    {
        keys.push_back(pair.first);
    }
    std::sort(keys.begin(), keys.end(), [](const std::wstring &a, const std::wstring &b)
              { return a.length() > b.length(); });

    while (i < len)
    {
        bool found = false;

        // 按长度从长到短尝试匹配
        for (const auto &key : keys)
        {
            if (i + key.length() <= len &&
                strs.compare(i, key.length(), key) == 0)
            {
                if (!loadCurrDictMat(key))
                    return false;
                // 找到匹配
                DictContent &content = currDict_->dicts_[key];
                vdict.push_back(content);
                i += key.length();
                found = true;
                break;
            }
        }

        if (!found)
        {
            return false;
        }
    }

    return true;
}
bool ImageHost::findDictTemplate(cv::Mat &binSource, cv::Mat &templateImg, double threshold, std::vector<cv::Rect> &rects)
{
    // 创建结果矩阵
    Mat result;
    int resultCols = binSource.cols - templateImg.cols + 1;
    int resultRows = binSource.rows - templateImg.rows + 1;
    if (resultCols <= 0 || resultRows <= 0)
        return false;
    result.create(resultRows, resultCols, CV_32FC1);

    // 执行模板匹配 [citation:1]
    matchTemplate(binSource, templateImg, result, TM_CCOEFF_NORMED);

    // 应用阈值过滤弱匹配 [citation:1]
    cv::threshold(result, result, threshold, 1., THRESH_TOZERO);

    // 查找所有超过阈值的匹配位置
    for (int y = 0; y < result.rows; y++)
    {
        for (int x = 0; x < result.cols; x++)
        {
            float matchValue = result.at<float>(y, x);
            if (matchValue >= threshold)
            {
                // 创建矩形并添加到结果向量
                cv::Rect matchRect(x, y, templateImg.cols, templateImg.rows);
                rects.push_back(matchRect);
            }
        }
    }

    // 返回是否找到匹配结果
    return !rects.empty();
}
cv::Rect ImageHost::matchDictTemplate(cv::Mat &binSource, const std::vector<DictContent> &vdict, double threshold)
{
    // 创建结果矩阵
    Mat result;
    int resultCols = binSource.cols;
    int resultRows = binSource.rows;
    result.create(resultRows, resultCols, CV_32FC1);
    for (int i = 0; i < vdict.size(); i++)
    {
        matchTemplate(binSource, vdict[i].mat_, result, TM_CCOEFF_NORMED);

        cv::threshold(result, result, threshold, 1., THRESH_TOZERO);
    }
    return Rect(-1, -1, -1, -1);
}
// 使用掩码的模板匹配
Rect ImageHost::findTemplateWithMask(Mat &source, Mat &templateImg, double threshold)
{
    // 分离模板图像的通道
    vector<Mat> templateChannels;
    split(templateImg, templateChannels);

    Mat templateMask;
    // 如果有4个通道，第4个是Alpha通道
    if (templateImg.channels() == 4)
    {
        templateMask = templateChannels[3]; // Alpha通道作为掩码
    }
    else
    {
        // 如果没有Alpha通道，创建全白掩码
        templateMask = Mat::ones(templateImg.size(), CV_8UC1) * 255;
    }

    // 将模板转换为3通道（移除Alpha）
    Mat templateRGB, grayTemplate;
    if (templateImg.channels() == 4)
    {
        vector<Mat> rgbChannels(templateChannels.begin(), templateChannels.begin() + 3);
        merge(rgbChannels, templateRGB);
    }
    else
    {
        templateRGB = templateImg.clone();
    }
    cvtColor(templateRGB, grayTemplate, COLOR_BGR2GRAY);

    cv::Mat graySource;
    if (source.channels() == 3)
    {
        cvtColor(source, graySource, COLOR_BGR2GRAY);
    }
    else if (source.channels() == 4)
    {
        cvtColor(source, graySource, COLOR_BGRA2GRAY);
    }
    else
    {
        graySource = source;
    }

    // 创建结果矩阵
    Mat result;
    int resultCols = source.cols - templateRGB.cols + 1;
    int resultRows = source.rows - templateRGB.rows + 1;
    result.create(resultRows, resultCols, CV_32FC1);

    // 使用掩码进行模板匹配
    matchTemplate(graySource, grayTemplate, result, TM_CCOEFF_NORMED, templateMask);

    // 应用阈值
    cv::threshold(result, result, threshold, 1., THRESH_TOZERO);

    // 查找最佳匹配
    double minVal, maxVal;
    Point minLoc, maxLoc;
    minMaxLoc(result, &minVal, &maxVal, &minLoc, &maxLoc);

    if (maxVal >= threshold)
    {
        return Rect(maxLoc.x, maxLoc.y, templateRGB.cols, templateRGB.rows);
    }

    return Rect(-1, -1, -1, -1);
}
cv::Rect ImageHost::findTemplate(cv::Mat &source, cv::Mat &templateImg, double threshold)
{
    // 转换为灰度图以减少计算量 [citation:1]
    cv::Mat graySource, grayTemplate;
    if (source.channels() == 3)
    {
        cvtColor(source, graySource, COLOR_BGR2GRAY);
    }
    else if (source.channels() == 4)
    {
        cvtColor(source, graySource, COLOR_BGRA2GRAY);
    }
    else
    {
        graySource = source;
    }
    if (templateImg.channels() == 3)
    {
        cvtColor(templateImg, grayTemplate, COLOR_BGR2GRAY);
    }
    else if (templateImg.channels() == 4)
    {
        cvtColor(templateImg, grayTemplate, COLOR_BGRA2GRAY);
    }
    else
    {
        grayTemplate = templateImg;
    }

    // 创建结果矩阵
    Mat result;
    int resultCols = source.cols - templateImg.cols + 1;
    int resultRows = source.rows - templateImg.rows + 1;
    if (resultCols <= 0 || resultRows <= 0)
    {
        return Rect(-1, -1, -1, -1);
    }
    result.create(resultRows, resultCols, CV_32FC1);

    // 执行模板匹配 [citation:1]
    matchTemplate(graySource, grayTemplate, result, TM_CCOEFF_NORMED);

    // 应用阈值过滤弱匹配 [citation:1]
    cv::threshold(result, result, threshold, 1., THRESH_TOZERO);

    // 查找最佳匹配位置
    double minVal, maxVal;
    Point minLoc, maxLoc;
    minMaxLoc(result, &minVal, &maxVal, &minLoc, &maxLoc);

    if (maxVal >= threshold)
    {
        return Rect(maxLoc.x, maxLoc.y, templateImg.cols, templateImg.rows);
    }

    return Rect(-1, -1, -1, -1);
}
